package org.example.core.processor;

import javassist.*;
import javassist.bytecode.*;
import org.example.core.template.CodeTemplate;
import org.example.core.util.ReflectUtils;
import org.example.core.util.AnnotationStoreUtils;
import org.springframework.aop.aspectj.annotation.AnnotationAwareAspectJAutoProxyCreator;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.beans.factory.config.*;
import org.springframework.beans.factory.support.*;
import org.springframework.context.annotation.AutoProxyRegistrar;
import org.springframework.core.annotation.MergedAnnotations;
import org.springframework.core.io.Resource;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.core.type.MethodMetadata;
import org.springframework.core.type.classreading.MetadataReader;
import org.springframework.core.type.classreading.SimpleMetadataReaderFactory;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Component;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import java.lang.annotation.*;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

@Component
public class TransactionalBeanFactoryPostProcessor implements BeanFactoryPostProcessor {

    private static final String TRANSACTIONAL_ANNOTATION_FIELD_PREFIX = "dynamic$$transactional$$annotation$$";

    private static final String INJECTION_SETTER_METHOD_PREFIX = "injection$$setter$$";

    private static final String PLATFORM_TRANSACTION_MANAGER_FIELD_NAME = "dynamic$$platformTransactionManager";

    /**
     * 类生成计数器
     */
    private static final AtomicInteger CLASS_COUNTER = new AtomicInteger();

    @Override
    public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
        try {
            // 拿到BeanDefinitionRegistry
            // 因为要涉及到BeanDefinition的替换, 需要removeBeanDefinition、registerBeanDefinition的能力
            BeanDefinitionRegistry registry = (BeanDefinitionRegistry) beanFactory;

            // 移除掉原本的Transactional功能
            // 这里仅仅是处理了EnableTransactionalManagement(AdviceMode mode() default AdviceMode.PROXY == PROXY)的情况
            // AdviceMode.ASPECTJ的情况可以自行去处理, 不过这个也不常用
            for (String beanDefinitionName : registry.getBeanDefinitionNames()) {
                BeanDefinition beanDefinition = registry.getBeanDefinition(beanDefinitionName);
                String beanClassName = beanDefinition.getBeanClassName();
                if (Objects.equals(beanClassName, AnnotationAwareAspectJAutoProxyCreator.class.getCanonicalName())) {
                    registry.removeBeanDefinition(beanDefinitionName);
                    break;
                }
            }

            // 通过ASM去读取类的元数据信息, 图省事也可以直接将beanClassName解析为beanClass
            // 能避免类被加载到JVM就避免一下
            SimpleMetadataReaderFactory metadataReaderFactory = new SimpleMetadataReaderFactory();

            for (String beanDefinitionName : registry.getBeanDefinitionNames()) {
                BeanDefinition beanDefinition = registry.getBeanDefinition(beanDefinitionName);

                // 只要AbstractBeanDefinition的子类(所有BeanDefinition都是它的子类), 需要使用到其中的cloneBeanDefinition
                if (beanDefinition instanceof AbstractBeanDefinition abd) {
                    // 因为Spring并没有真正的拿到beanClass实例, 仅存在beanClassName
                    // 需要通过beanClassName来操作并执行筛选, 最终拿到beanClass
                    String beanClassName = abd.getBeanClassName();
                    if (beanClassName == null) {
                        continue;
                    }

                    // 校验beanDefinition的resource即它的源文件是否存在
                    // 如果不存在大概率是代码内动态生成的beanDefinition, 例如CGLIB, 此时需要跳过 (因为无法使用MetadataReader读取)
                    Resource resource = abd.getResource();
                    if (resource == null || !resource.exists()) {
                        continue;
                    }

                    // 通过ASM读取Metadata
                    MetadataReader metadataReader = metadataReaderFactory.getMetadataReader(beanClassName);
                    AnnotationMetadata annotationMetadata = metadataReader.getAnnotationMetadata();

                    // 检查是否有需要重写的方法, 至少存在一个才会创建子类
                    Set<String> overwriteMethodNames = getOverwriteMethodNames(annotationMetadata);
                    if (CollectionUtils.isEmpty(overwriteMethodNames)) {
                        continue;
                    }

                    // 利用javassist创建子类
                    ClassPool cp = ClassPool.getDefault();
                    CtClass parent = cp.get(beanClassName);
                    CtClass subclass = cp.makeClass(generateSubclassName(beanClassName), parent);

                    // 通过依赖注入获取PlatformTransactionManager
                    addAutowireField(subclass, PlatformTransactionManager.class.getCanonicalName(), PLATFORM_TRANSACTION_MANAGER_FIELD_NAME);

                    // 重写成员
                    overwriteFields(subclass, parent);
                    overwriteMethods(subclass, parent, overwriteMethodNames);
                    overwriteConstructors(subclass, parent);

                    // 创建新的BeanDefinition
                    AbstractBeanDefinition newBeanDefinition = abd.cloneBeanDefinition();
                    Class<?> aClass = subclass.toClass();
                    newBeanDefinition.setBeanClass(aClass);
                    newBeanDefinition.setBeanClassName(newBeanDefinition.getBeanClass().getCanonicalName());

                    replaceBeanDefinition(registry, beanDefinitionName, newBeanDefinition);
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    private void replaceBeanDefinition(BeanDefinitionRegistry registry, String beanDefinitionName, BeanDefinition beanDefinition) {
        if (registry.containsBeanDefinition(beanDefinitionName)) {
            registry.removeBeanDefinition(beanDefinitionName);
        }
        registry.registerBeanDefinition(beanDefinitionName, beanDefinition);
    }

    /**
     * 添加一个带有依赖注入注解的字段
     *
     * @param subclass 动态生成的子类
     * @param type 字段数据类型全限定名称
     * @param name 字段名称
     */
    private void addAutowireField(CtClass subclass, String type, String name) throws CannotCompileException {
        ConstPool cp = subclass.getClassFile().getConstPool();
        CtField field = CtField.make("public " + type + " " + name + ";", subclass);
        AnnotationsAttribute annotationsAttribute = new AnnotationsAttribute(cp, AnnotationsAttribute.visibleTag);
        javassist.bytecode.annotation.Annotation annotation = new javassist.bytecode.annotation.Annotation(Autowired.class.getCanonicalName(), cp);
        annotationsAttribute.addAnnotation(annotation);
        field.getFieldInfo().addAttribute(annotationsAttribute);
        subclass.addField(field);
    }

    private void overwriteFields(CtClass subclass, CtClass parent) throws ClassNotFoundException, CannotCompileException, NotFoundException {
        for (CtField field : parent.getDeclaredFields()) {
            int modifiers = field.getModifiers();
            if (Modifier.isStatic(modifiers) || Modifier.isFinal(modifiers) || Modifier.isAbstract(modifiers)) {
                continue;
            }
            // 对于Object类字段以及无注解的字段均跳过
            String declaringClassName = field.getDeclaringClass().getName();
            if (declaringClassName.equals(Object.class.getCanonicalName()) || field.getAnnotations().length == 0) {
                continue;
            }
            // 私有属性时单独处理, 为了兼容私有属性Autowired、Value, 子类动态生成set方法, 通过反射强制设置到父类
            // 这里暂时没考虑Resource, Inject这类注解, 可以自己加上
            if (Modifier.isPrivate(modifiers)) {
                Object annotation = field.getAnnotation(Autowired.class);
                if (annotation == null) {
                    annotation = field.getAnnotation(Value.class);
                }
                if (annotation != null) {
                    CtMethod setter = CtNewMethod.make(
                            CtClass.voidType,
                            INJECTION_SETTER_METHOD_PREFIX + field.getName(),
                            new CtClass[]{field.getType()},
                            new CtClass[0],
                            CodeTemplate.generateInjectionSetterCode(field.getName()),
                            subclass
                    );
                    // 将字段上的注解标记到方法上, 只要可被标记在方法之上的注解
                    AnnotationsAttribute annotationsAttribute = new AnnotationsAttribute(subclass.getClassFile().getConstPool(), AnnotationsAttribute.visibleTag);
                    for (Object fieldAnnotation : field.getAnnotations()) {
                        Annotation javaAnnotation = (Annotation) fieldAnnotation;
                        Target target = javaAnnotation.annotationType().getAnnotation(Target.class);
                        boolean candidate = false;
                        if (target == null) {
                            candidate = true;
                        } else {
                            for (ElementType elementType : target.value()) {
                                if (elementType == ElementType.METHOD) {
                                    candidate = true;
                                    break;
                                }
                            }
                        }
                        if (candidate) {
                            annotationsAttribute.addAnnotation(ReflectUtils.toJavassistAnnotation(javaAnnotation, subclass.getClassFile().getConstPool()));
                        }
                    }
                    if (annotationsAttribute.numAnnotations() > 0) {
                        setter.getMethodInfo().addAttribute(annotationsAttribute);
                    }
                    subclass.addMethod(setter);
                }
            } else {
                // 非私有属性直接重写
                subclass.addField(new CtField(field, subclass));
            }
        }
    }

    private void overwriteMethods(CtClass subclass, CtClass parent, Set<String> overwriteMethodNames) throws NotFoundException, CannotCompileException, ClassNotFoundException {
        // 遍历方法, 为每个标记Transactional注解的方法都创建重写, 并保存它们的注解实例到字段中
        for (CtMethod method : parent.getMethods()) {
            if (!overwriteMethodNames.contains(method.getName())) {
                continue;
            }
            // 避免无注解的同名重载方法也开启事务
            Object annotation = method.getAnnotation(Transactional.class);
            if (annotation == null) {
                continue;
            }
            // 1. 将注解实例存储到类中, 方便操作事务时获取相关配置
            // 2. 将添加的属性名以及值put到map中, 在后续通过构造器取出存放到字段中
            Transactional transactional = (Transactional) annotation;
            String annotationFieldName = TRANSACTIONAL_ANNOTATION_FIELD_PREFIX + method.getName();
            subclass.addField(CtField.make("private " + Transactional.class.getCanonicalName() + " " + annotationFieldName + ";", subclass));
            AnnotationStoreUtils.putValue(subclass.getName(), annotationFieldName, transactional);

            // 添加重写方法, 植入事务控制代码
            String methodName = method.getName();
            CtClass[] parameterTypes = method.getParameterTypes();
            int parameterCount = parameterTypes.length;
            CtClass returnType = method.getReturnType();
            String returnTypeString = returnType.equals(CtClass.voidType) ? null : returnType.getName();
            CtMethod ctMethod = CtNewMethod.make(
                    returnType,
                    methodName,
                    parameterTypes,
                    method.getExceptionTypes(),
                    CodeTemplate.generateTransactionControlCode(
                            annotationFieldName,
                            PLATFORM_TRANSACTION_MANAGER_FIELD_NAME,
                            returnTypeString,
                            methodName,
                            parameterCount
                    ),
                    subclass
            );

            subclass.addMethod(ctMethod);
        }
    }

    private void overwriteConstructors(CtClass subclass, CtClass parent) throws CannotCompileException, NotFoundException {
        for (CtConstructor declaredConstructor : parent.getDeclaredConstructors()) {
            CtConstructor constructor = new CtConstructor(declaredConstructor, subclass, new ClassMap());
            constructor.setBody(CodeTemplate.generateConstructorCode(constructor.getParameterTypes().length));
            subclass.addConstructor(constructor);
        }
    }
    private String generateSubclassName(String originalClassName) {
        return String.format("%s$$dynamic$$%d", originalClassName, CLASS_COUNTER.getAndIncrement());
    }

    /**
     * 通过Metadata信息读取所有需要重写的方法名称
     *
     * @param annotationMetadata AnnotationMetadata
     * @return nameSet
     */
    @Nullable
    private Set<String> getOverwriteMethodNames(AnnotationMetadata annotationMetadata) {
        Set<String> nameSet = null;
        for (MethodMetadata declaredMethod : annotationMetadata.getDeclaredMethods()) {
            MergedAnnotations annotations = declaredMethod.getAnnotations();
            if (annotations.isPresent(Transactional.class)) {
                if (nameSet == null) {
                    nameSet = new HashSet<>();
                }
                nameSet.add(declaredMethod.getMethodName());
            }
        }
        return nameSet;
    }
}
